import numpy as np
import torch.nn as nn
from datasets import *

d_model = 512  # 字 Embedding 的维度
d_ff = 2048  # 前向传播隐藏层维度
d_k = d_v = 64  # K(=Q), V的维度
n_layers = 6  # 有多少个encoder和decoder
n_heads = 8  # Multi-Head Attention设置为8


class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        # 其中，i是向量中每个维度的下标，从0到d_model-1；^表示幂运算。这个公式中的10000是一个常数，
        # 为了确保位置编码矩阵中不同位置的编码能够有所不同，从而区分不同位置的输入。根据公式，随着i的增大，位置编码矩阵中的每个维度都会被分配到具有更高频率的正弦/余弦函数中。
        # 这样可以确保不同位置的编码之间有足够的差异，从而提供位置信息。
        #
        # 然而，当pos=0时，由于分母会变为0，会导致公式的结果出现异常。
        # 因此，代码使用if-else语句对这种情况进行了特判：当pos=0时，将位置向量的所有维度设为0。这样，位置编码矩阵的第一行就是一个全零向量，表示序列的起始位置。
        pos_table = np.array([
            [pos / np.power(10000, 2 * i / d_model) for i in range(d_model)]
            if pos != 0 else np.zeros(d_model) for pos in range(max_len)])
        # 根据公式分别计算
        # 0::2 双冒号后的表示步长
        pos_table[1:, 0::2] = np.sin(pos_table[1:, 0::2])  # 字嵌入维度为偶数时
        pos_table[1:, 1::2] = np.cos(pos_table[1:, 1::2])  # 字嵌入维度为奇数时
        self.pos_table = torch.FloatTensor(pos_table).cuda()  # enc_inputs: [seq_len, d_model]

    def forward(self, enc_inputs):  # enc_inputs: [batch_size, seq_len, d_model]
        # 将 pos_table 加到 enc_inputs 的前 enc_inputs.size(1) 行，即只对当前序列的每个位置进行位置编码(为输入序列的前seq_length个位置添加唯一的位置编码)
        enc_inputs += self.pos_table[:enc_inputs.size(1), :]
        # 用于在神经网络中实现 dropout 操作，即在训练过程中随机将一些神经元的输出值置为 0，以避免过拟合。
        # 它的输入是一个张量和一个 dropout 概率，输出是一个与输入相同形状的张量，其中随机置 0 的元素的概率为输入的 dropout 概率。它通常用于在全连接层或卷积层的输出之后执行。
        return self.dropout(enc_inputs.cuda())


# Mask掉停用词
# 生成注意力掩码矩阵
# NOTE: 在进行注意力机制计算时将填充位置的信息屏蔽掉，防止填充位置的信息影响模型的计算和训练。
# WHY: 在输入序列中，为了使所有序列长度一致，通常需要对长度不足的序列进行填充。
#  而对于这些填充的位置，实际上是不需要进行处理的，因为它们不包含任何有效信息。在进行自注意力机制计算时，padding mask被用来指示模型哪些位置是填充位置。
#  具体来说，如果一个位置是填充位置，那么它的padding mask值为1，否则为0。在进行softmax计算时，padding mask会使那些填充位置对应的softmax值趋近于0，从而将它们屏蔽掉，保证计算的准确性。
#  同样，当进行编码器-解码器注意力机制时，为了避免解码器中的位置访问编码器中的填充位置，需要对编码器输入序列进行padding mask处理。
#  在解码器的自注意力机制计算中，还需要额外加上subsequence mask，保证在解码器的每个位置上，只能使用当前位置之前的信息，避免信息泄漏。
def get_attn_pad_mask(seq_q, seq_k):
    # seq_q: [batch_size, seq_len] ,seq_k: [batch_size, seq_len]:表示batch_size个序列，每个序列有seq_len个单词
    batch_size, len_q = seq_q.size()
    batch_size, len_k = seq_k.size()
    # unsqueeze()是PyTorch中的一个函数，用于在张量（Tensor）中插入维度（dimension）[在指定位置上插入一个新的维度]。
    # 这个函数接受一个整数参数作为维度，它会在对应位置上增加一个大小为1的新维度。例如，对于一个形状为(3,4)的张量，调用unsqueeze(0)会返回一个形状为(1,3,4)的张量，而调用unsqueeze(1)则会返回一个形状为(3,1,4)的张量。
    # unsqueeze()函数通常用于构建新的张量，以便与其他张量进行广播（broadcasting），或者用于调整输入张量的形状以满足某些函数的输入要求。
    pad_attn_mask = seq_k.data.eq(0).unsqueeze(1)  # 判断 输入那些含有P(=0),用1标记 ,[batch_size, 1, len_k]
    # 函数返回一个形状为[batch_size, seq_len, seq_len]的张量，表示对于每个序列中的每个单词，NOTE: 哪些单词是padding（P）的需要被忽略的,即 将需要被忽略的位置标记为1
    return pad_attn_mask.expand(batch_size, len_q, len_k)  # 扩展成多维度


# Decoder 输入 Mask
# 生成self-attention的子序列掩码
# NOTE: 在自注意力机制中，为了避免当前位置对后续位置的影响，需要将当前位置后面的位置掩盖起来，不参与计算
#   因此，对于一个输入序列seq，函数生成一个上三角矩阵，对角线及其以上部分都是0，对角线以下部分都是1，即掩盖了当前位置之后的所有位置，返回掩码subsequence_mask
#   这里的mask属于 sequence mask: 防止未来信息泄露    https://www.zhihu.com/question/320615749
def get_attn_subsequence_mask(seq):  # seq: [batch_size, tgt_len]
    # 指定要生成的mask矩阵的shape
    attn_shape = [seq.size(0), seq.size(1), seq.size(1)]
    subsequence_mask = np.triu(np.ones(attn_shape), k=1)  # 生成上三角矩阵,[batch_size, tgt_len, tgt_len]
    # NOTE: tensor.byte()是一个函数，用于将张量的数据类型转换为8位无符号整数类型。它返回一个新的张量，其数据类型为torch.uint8，并且每个元素都在0到255之间
    subsequence_mask = torch.from_numpy(subsequence_mask).byte()  # [batch_size, tgt_len, tgt_len]
    return subsequence_mask


# scaled dot-product attention(缩放点积注意力)的实现
class ScaledDotProductAttention(nn.Module):
    def __init__(self):
        super(ScaledDotProductAttention, self).__init__()

    def forward(self, Q, K, V, attn_mask):
        # Q: [batch_size, n_heads, len_q, d_k] torch.Size([2, 8, 5, 64])
        # K: [batch_size, n_heads, len_k, d_k] torch.Size([2, 8, 5, 64])
        # V: [batch_size, n_heads, len_v(=len_k), d_v] torch.Size([2, 8, 5, 64])
        # attn_mask: [batch_size, n_heads, seq_len, seq_len]
        # NOTE:通过torch.matmul(Q, K.transpose(-1, -2))计算Q和K的点积得分，然后除以np.sqrt(d_k)作为缩放因子
        scores = torch.matmul(Q, K.transpose(-1, -2)) / np.sqrt(d_k)  # scores : [batch_size, n_heads, len_q, len_k] torch.Size([2, 8, 5, 5])
        # NOTE: 使用了一个掩码（mask），将无效位置的得分（score）设置为了一个非常小的负数，以避免这些无效位置的得分被计入最终的注意力分布中。
        #       具体的： attn_mask 是一个与 scores 的形状相同的布尔型张量，表示哪些位置需要被掩盖。
        #       在这行代码中，我们使用 masked_fill_ 函数将 attn_mask 中为 True 的位置对应的 scores 设置为一个非常小的负数 -1e9，以确保这些位置在 softmax 函数中被归一化后得到的权重接近于零，从而在最终的加权求和中不产生贡献。
        scores.masked_fill_(attn_mask, -1e9)  # 如果是停用词P就等于 0
        # 对scores进行softmax激活，得到了attn，这个张量用于对V进行加权求和
        attn = nn.Softmax(dim=-1)(scores)  # NOTE: dim=-1: 对最后一维单独做softmax attn:[batch_size, n_heads, len_q, len_k]
        # NOTE:
        #  在这个运算中，首先将 attn 和 V 做矩阵乘法，对于每个样本和每个注意力头，都会得到一个形状为 (seq_len, d_v) 的上下文张量
        #  相当于是对于每个[batch_size, n_heads]维度，[len_q, len_k]x[len_q, d_k]
        #  最终得到的上下文张量 context 的维度为 (batch_size, n_heads, seq_len, d_v)，
        #  其中 seq_len 表示序列长度，d_v 表示注意力头中向量的维度。可以看到，这个运算中只改变了最后一个维度，其他维度保持不变。
        context = torch.matmul(attn, V)  # context: [batch_size, n_heads, len_q, d_v]  torch.Size([2, 8, 5, 64])
        return context, attn


# 多头注意力机制模块(Encoder和Decoder公用)
class MultiHeadAttention(nn.Module):
    def __init__(self):
        super(MultiHeadAttention, self).__init__()
        # 这里初始化时W_Q,W_K,W_V的结构都是相同维度的线性层
        # 输入：词嵌入向量大小(输入数据的特征维度)   输出：d_k * n_heads
        # self-attention分为多头时满足： d_model = d_k * n_heads(相当于把向量维度分散成了n_heads个)
        self.W_Q = nn.Linear(d_model, d_k * n_heads, bias=False)
        self.W_K = nn.Linear(d_model, d_k * n_heads, bias=False)
        self.W_V = nn.Linear(d_model, d_v * n_heads, bias=False)
        self.fc = nn.Linear(n_heads * d_v, d_model, bias=False)

    def forward(self, input_Q, input_K, input_V, attn_mask):
        # input_Q: [batch_size, len_q, d_model]
        # input_K: [batch_size, len_k, d_model]
        # input_V: [batch_size, len_v(=len_k), d_model]
        # attn_mask: [batch_size, seq_len, seq_len]
        residual, batch_size = input_Q, input_Q.size(0)
        # 1. 这里先通过input经过线性层得到 Q,K,V
        # Q、K、V是通过输入的Q、K、V矩阵分别乘以三个可训练的权重矩阵W_Q、W_K、W_V得到的，因此可以学习到输入的不同表示。

        # 2. 将Q、K、V矩阵分别进行reshape和transpose变换，
        # 得到Q、K、V矩阵的形状为[batch_size, n_heads, len_q, d_k]、[batch_size, n_heads, len_k, d_k]、[batch_size, n_heads, len_v(=len_k), d_v]
        # 其中n_heads表示头数，即多头注意力机制的个数；d_k和d_v表示分头注意力机制中每个头的特征维度，即d_model/n_heads。
        # 在这里将512的维度切分成8个64维度的矩阵
        # NOTE:
        #   在标准的 Transformer 模型中，len_q、len_k、len_v 的值是相等的，且通常设置为一个较小的值，以便使模型能够并行处理多个注意力头的计算。
        #   因此，在标准的 Transformer 模型中，输入序列长度与 len_q、len_k、len_v 的关系可以简单地表示为：
        #   len_q = len_k = len_v = L   (其中，L 表示每个注意力头中参与计算的序列长度)
        Q = self.W_Q(input_Q).view(batch_size, -1, n_heads, d_k).transpose(1, 2)  # Q: [batch_size, n_heads, len_q, d_k]
        K = self.W_K(input_K).view(batch_size, -1, n_heads, d_k).transpose(1, 2)  # K: [batch_size, n_heads, len_k, d_k]
        V = self.W_V(input_V).view(batch_size, -1, n_heads, d_v).transpose(1, 2)  # V: [batch_size, n_heads, len_v(=len_k), d_v]
        # NOTE: 通过repeat()函数的操作，我们将一个(batch_size, seq_len, seq_len)的张量复制了n_heads次，并沿着第2个维度(n_heads维度)重复了seq_len次，以便与其他张量进行并行计算。
        # 这样就得到了一个形状为(batch_size, n_heads, seq_len, seq_len)的张量，其中每个头都具有相同的attention mask，可以与其他头一起进行并行计算。
        attn_mask = attn_mask.unsqueeze(1).repeat(1, n_heads, 1, 1)     # attn_mask : [batch_size, n_heads, seq_len, seq_len]
        context, attn = ScaledDotProductAttention()(Q, K, V, attn_mask)
        # context: [batch_size, n_heads, len_q, d_v]  NOTE: 缩放点积注意力后的张量大小与V的大小相同
        # NOTE: 一般情况下，通过将d_k=d_q=d_v，我们可以确保在每个头中使用的查询向量、键向量和值向量的维度相同。这个特定的维度选择方式可以使不同头之间的注意力计算更具可比性，从而使模型更容易学习
        # note:输出2：attn: [batch_size, n_heads, len_q, len_k]
        # 对输入数据的每个头进行Scaled Dot-Product Attention计算，得到n_heads个context矩阵，将这些矩阵进行concatenate和reshape操作，得到形状为[batch_size, len_q, n_heads * d_v]的context矩阵
        context = context.transpose(1, 2).reshape(batch_size, -1, n_heads * d_v)
        # context: [batch_size, len_q, n_heads * d_v] == [batch_size, len_q, d_model]
        # 对context矩阵进行一个全连接层变换，得到输出矩阵
        output = self.fc(context)  # note:输出1：[batch_size, len_q, d_model]
        # 使用残差连接（residual connection）和层归一化（layer normalization）操作得到最终的输出结果和注意力权重（attn）
        return nn.LayerNorm(d_model).cuda()(output + residual), attn


# 前馈神经网络(Encoder和Decoder公用)
# NOTE: 这里前馈神经网络的计算过程：
#       FFN(h_i) = max(0, h_i W_1 + b_1) W_2 + b_2
#       这个序列模型的作用是对每个位置的隐藏表示进行一些非线性变换，以增强位置之间的交互和信息流动
# NOTE: 前馈层不改变输入的结构
class PoswiseFeedForwardNet(nn.Module):
    def __init__(self):
        super(PoswiseFeedForwardNet, self).__init__()
        # NOTE: nn.Sequential 是 PyTorch 中一个用于构建序列模型的类。它可以将一组层组合成一个网络，形成一个前向传播函数，从而方便地进行模型的定义和训练。
        #       在使用 nn.Sequential 构建模型时，可以按照顺序将不同的层添加到序列中，每个层的输出都作为下一个层的输入。
        self.fc = nn.Sequential(
            nn.Linear(d_model, d_ff, bias=False),
            nn.ReLU(),
            nn.Linear(d_ff, d_model, bias=False))

    def forward(self, inputs):  # inputs: [batch_size, seq_len, d_model]
        residual = inputs
        output = self.fc(inputs)
        # 前馈中包含了 残差和层归一化
        return nn.LayerNorm(d_model).cuda()(output + residual)  # [batch_size, seq_len, d_model]


# Transformer模型中Encoder的基本组成单元——Encoder Layer
# Encoder Layer包含了一个多头注意力机制（Multi-Head Attention）和一个前馈神经网络（Poswise Feed Forward Net）
class EncoderLayer(nn.Module):
    def __init__(self):
        super(EncoderLayer, self).__init__()
        self.enc_self_attn = MultiHeadAttention()  # 多头注意力机制
        self.pos_ffn = PoswiseFeedForwardNet()  # 前馈神经网络

    # 在Encoder Layer的forward函数中，输入enc_inputs表示Encoder的输入，它的形状为[batch_size, src_len, d_model]，
    # 其中batch_size表示一批数据的大小，src_len表示输入序列的长度(seq_len)，d_model表示模型的维度(字embedding维度)。
    # enc_self_attn_mask是一个掩码矩阵，用于在多头注意力机制中屏蔽不需要关注的位置
    def forward(self, enc_inputs, enc_self_attn_mask):  # enc_inputs: [batch_size, src_len, d_model]
        # 输入3个enc_inputs分别与W_q、W_k、W_v相乘得到Q、K、V             # enc_self_attn_mask: [batch_size, src_len, src_len]
        enc_outputs, attn = self.enc_self_attn(enc_inputs, enc_inputs, enc_inputs, enc_self_attn_mask)
        # NOTE: 自注意力的输出也就是：enc_outputs = context(注意力上下文张量经过残差以及层归一化后的结果)
        # enc_outputs: [batch_size, src_len, d_model]
        # attn: [batch_size, n_heads, src_len, src_len]
        # 前馈后的残差和BL加在了前馈网络中
        enc_outputs = self.pos_ffn(enc_outputs)
        # NOTE: src_len在这里表示输入序列的长度
        # enc_outputs: [batch_size, src_len, d_model]
        return enc_outputs, attn


# 编码器模块
class Encoder(nn.Module):
    def __init__(self):
        super(Encoder, self).__init__()
        self.src_emb = nn.Embedding(src_vocab_size, d_model)  # 定义词嵌入层的维度
        self.pos_emb = PositionalEncoding(d_model)  # 加入位置信息
        # nn.ModuleList是PyTorch中的一个容器，它可以容纳一个nn.Module对象列表。它类似于普通的Python列表，但专门设计用于容纳PyTorch模块。
        self.layers = nn.ModuleList([EncoderLayer() for _ in range(n_layers)])

    def forward(self, enc_inputs):  # enc_inputs: [batch_size, src_len]  这里的src_len其实就是 seq_len
        enc_outputs = self.src_emb(enc_inputs)  # enc_outputs: [batch_size, src_len, d_model]
        enc_outputs = self.pos_emb(enc_outputs)  # enc_outputs: [batch_size, src_len, d_model]  [批量大小，句子序列的长度，字向量维度]
        # 使用get_attn_pad_mask函数生成了一个enc_self_attn_mask张量，维度为[batch_size, src_len, src_len]
        # NOTE: 表示对源语言句子的mask，避免句子中的padding token对注意力权重的影响
        enc_self_attn_mask = get_attn_pad_mask(enc_inputs,enc_inputs)   # enc_self_attn_mask: [batch_size, src_len, src_len]
        enc_self_attns = []
        for layer in self.layers:
            enc_outputs, enc_self_attn = layer(enc_outputs,enc_self_attn_mask)
            # enc_outputs :   [batch_size, src_len, d_model]
            # enc_self_attn : [batch_size, n_heads, src_len, src_len]
            enc_self_attns.append(enc_self_attn)
            # 这里返回最后一层EncoderLayer的输出 和 每一层的注意力张量
        return enc_outputs, enc_self_attns


# 解码器层模块
# 解码器由多个解码器Layer连接组成
class DecoderLayer(nn.Module):
    def __init__(self):
        super(DecoderLayer, self).__init__()
        self.dec_self_attn = MultiHeadAttention()
        self.dec_enc_attn = MultiHeadAttention()
        self.pos_ffn = PoswiseFeedForwardNet()

    # 解码器
    def forward(self, dec_inputs, enc_outputs, dec_self_attn_mask, dec_enc_attn_mask):
        # dec_inputs: [batch_size, tgt_len, d_model]
        # enc_outputs: [batch_size, src_len, d_model]
        # dec_self_attn_mask: [batch_size, tgt_len, tgt_len]
        # dec_enc_attn_mask: [batch_size, tgt_len, src_len]

        # 每个Decoder_layer的结构： 解码器输入掩码多头注意力 + 编码器输出多头注意力 + 前馈
        # NOTE: Decoder_layer 中的第一个多头注意力：Masked Multi-Head Attention:input(Q,K,V):dec_inputs
        dec_outputs, dec_self_attn = self.dec_self_attn(dec_inputs, dec_inputs, dec_inputs, dec_self_attn_mask)
        # dec_outputs: [batch_size, tgt_len, d_model]
        # dec_self_attn: [batch_size, n_heads, tgt_len, tgt_len]

        # Decoder_layer 中的第二个多头注意力：input(Q,K,V):
        # NOTE: Decoder 的 Multi-Head Attention 的输入来自两部分，K，V 矩阵来自Encoder的输出(enc_outputs)，Q矩阵来自 Masked Multi-Head Attention的输出(dec_outputs)
        dec_outputs, dec_enc_attn = self.dec_enc_attn(dec_outputs, enc_outputs, enc_outputs, dec_enc_attn_mask)
        # NOTE: 注意这里的输出维度与实际的句子长度信息之间的对应关系：
        #       dec_outputs: [batch_size, tgt_len, d_model] <-- [batch_size, len_q, d_model]
        #       dec_enc_attn: [batch_size, h_heads, tgt_len, src_len]  <-- [batch_size,h_heads,len_q,len_k]
        dec_outputs = self.pos_ffn(dec_outputs)  # dec_outputs: [batch_size, tgt_len, d_model]
        # 返回： 解码器输出，解码器自注意力矩阵，解码器编码器注意力矩阵
        return dec_outputs, dec_self_attn, dec_enc_attn


# Decoder解码器类
class Decoder(nn.Module):
    def __init__(self):
        super(Decoder, self).__init__()
        self.tgt_emb = nn.Embedding(tgt_vocab_size, d_model)
        self.pos_emb = PositionalEncoding(d_model)
        self.layers = nn.ModuleList([DecoderLayer() for _ in range(n_layers)])

    def forward(self, dec_inputs, enc_inputs, enc_outputs):
        # NOTE： 目标语言的输入序列：dec_inputs: [batch_size, tgt_len]
        #       源语言的输入序列：enc_intpus: [batch_size, src_len]
        #       经过Encoder计算后得到的编码结果： enc_outputs: [batsh_size, src_len, d_model]
        dec_outputs = self.tgt_emb(dec_inputs)  # [batch_size, tgt_len, d_model]
        dec_outputs = self.pos_emb(dec_outputs).cuda()  # [batch_size, tgt_len, d_model]
        # NOTE:通过get_attn_pad_mask和get_attn_subsequence_mask得到解码器自注意力的mask，这些mask用于遮盖输入序列中的padding符号和未来的信息
        dec_self_attn_pad_mask = get_attn_pad_mask(dec_inputs, dec_inputs).cuda()  # [batch_size, tgt_len, tgt_len]
        dec_self_attn_subsequence_mask = get_attn_subsequence_mask(dec_inputs).cuda()  # [batch_size, tgt_len, tgt_len]
        # NOTE: 生成Decoder自注意力的mask矩阵:
        #   dec_self_attn_pad_mask 和 dec_self_attn_subsequence_mask 首先被相加，得到的结果会存在两种情况：
        #   如果某个位置同时在两个矩阵中都被标记为0，那么相加后它的值为0；否则它的值为1
        #   使用torch.gt()函数将所有值大于0的位置标记为1，所有其他位置标记为0，从而生成一个新的mask矩阵dec_self_attn_mask。（torch.gt()要求两个张量大小相同，否则需要将形状较小的张量扩展后再比较. 例如：先unsqueeze）
        #   最后使用.cuda()将该张量移动到GPU上。最终生成的dec_self_attn_mask矩阵是一个三维矩阵，其维度为[batch_size, tgt_len, tgt_len]，用于在Decoder的自注意力中屏蔽掉无效的信息
        dec_self_attn_mask = torch.gt((dec_self_attn_pad_mask + dec_self_attn_subsequence_mask), 0).cuda()  # [batch_size, tgt_len, tgt_len]
        # 通过get_attn_pad_mask计算解码器和编码器之间的注意力的mask，这个mask同样用于遮盖输入序列中的padding符号
        dec_enc_attn_mask = get_attn_pad_mask(dec_inputs, enc_inputs)  # [batc_size, tgt_len, src_len]
        dec_self_attns, dec_enc_attns = [], []
        for layer in self.layers:
            # dec_outputs: [batch_size, tgt_len, d_model]
            # dec_self_attn: [batch_size, n_heads, tgt_len, tgt_len]
            # dec_enc_attn: [batch_size, h_heads, tgt_len, src_len]
            dec_outputs, dec_self_attn, dec_enc_attn = layer(dec_outputs, enc_outputs, dec_self_attn_mask, dec_enc_attn_mask)
            dec_self_attns.append(dec_self_attn)
            dec_enc_attns.append(dec_enc_attn)
        return dec_outputs, dec_self_attns, dec_enc_attns


# Transformer模型结构
class Transformer(nn.Module):
    def __init__(self):
        super(Transformer, self).__init__()
        # NOTE:调用cuda()方法时，PyTorch会将模型中所有的可训练参数和缓存（如中间变量）从CPU内存转移到GPU内存，并将模型的计算操作设置为在GPU上执行。这样就可以利用GPU的并行计算能力来加速模型的训练和推理过程
        # 如果我们需要在CPU上使用模型，可以通过调用cpu()方法将模型转回CPU内存
        self.Encoder = Encoder().cuda()
        self.Decoder = Decoder().cuda()
        # NOTE: 线性层，用于将 Decoder 输出的 d_model 维向量映射到目标词表的大小
        self.projection = nn.Linear(d_model, tgt_vocab_size, bias=False).cuda()

    def forward(self, enc_inputs, dec_inputs):
        # enc_inputs: [batch_size, src_len]:源语言句子
        # dec_inputs: [batch_size, tgt_len]:目标语言句子
        # NOTE: transformer在训练和预测过程是的输入是不一样的
        #       在训练过程中，模型的输入通常包含源语言句子和目标语言句子，其中源语言句子是编码器的输入，目标语言句子是解码器的输入。
        #   目标语言句子还需要在每个位置上移动一个位置，作为解码器的输出，即模型需要预测目标语言句子中每个位置的下一个词。
        #       在预测过程中，模型的输入通常只包含源语言句子，即需要将源语言句子作为编码器的输入，然后通过解码器生成目标语言句子。
        #   生成目标语言句子的过程是一个逐步生成的过程，即在每一步生成一个词，然后将这个词作为下一步生成的输入，直到生成完整的目标语言句子。
        #       需要注意的是，在预测过程中，模型需要不断地进行解码器的前向计算，并生成每个位置的输出，因此相比训练过程，预测过程需要更多的计算资源。
        #   同时，在预测过程中，为了防止生成重复的词或者无限循环，还需要进行一些特殊的处理，例如采用 Beam Search 算法等。
        enc_outputs, enc_self_attns = self.Encoder(enc_inputs)
        # enc_outputs: [batch_size, src_len, d_model],
        # enc_self_attns: [n_layers, batch_size, n_heads, src_len, src_len]
        # NOTE: 这里的Decoder中的enc_inputs的输入是用来计算 decode_encode_mask
        dec_outputs, dec_self_attns, dec_enc_attns = self.Decoder(dec_inputs, enc_inputs, enc_outputs)
        # dec_outpus    : [batch_size, tgt_len, d_model],
        # NOTE:这里的 dec_self_attns 和 dec_enc_attn 都是存储了注意力数据的list
        # dec_self_attns: [n_layers, batch_size, n_heads, tgt_len, tgt_len],
        # dec_enc_attn  : [n_layers, batch_size, n_heads, tgt_len, src_len]
        dec_logits = self.projection(dec_outputs)
        # IMPO: 最终的输出维度：dec_logits: [batch_size, tgt_len, tgt_vocab_size]
        # NOTE: Transformer最终的输出是经过解码器输出层计算得到的概率分布，表示每个位置对应的词汇表中各个单词的概率。
        #       一般来说，最终的输出是一个形状为(batch_size, seq_len, vocab_size)的张量，其中每个元素(i,j,k)表示第i个样本在位置j处输出单词k的概率。
        #        在训练过程中，我们可以使用这个概率分布计算交叉熵损失函数，并进行反向传播更新模型参数。
        #        在推理过程中，我们可以使用这个概率分布来生成文本，例如采用贪心策略或者采样策略选取概率最高的单词作为输出。
        # note: 通过调用 view 方法将其变形为 [batch_size * tgt_len, tgt_vocab_size] 的形状，方便之后计算损失
        # 同时也返回了 Encoder 自注意力矩阵、Decoder 自注意力矩阵和Decoder-Encoder 注意力矩阵，用于可视化和调试
        return dec_logits.view(-1, dec_logits.size(-1)), enc_self_attns, dec_self_attns, dec_enc_attns
